package com.example.demo.demos.rest;

import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Configuration;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicInteger;

@Configuration
public class ToolAutoReg implements BeanPostProcessor, ApplicationContextAware {

    private AtomicInteger c = new AtomicInteger(0);

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        // 获取bean的所有方法
        Method[] methods = bean.getClass().getDeclaredMethods();

        for (Method method : methods) {
            // 检查方法是否有@RequestMapping注解
            if (method.isAnnotationPresent(Tool.class)) {
                BeanDefinitionBuilder newBean = BeanDefinitionBuilder.genericBeanDefinition(MethodToolCallbackProvider.class).addConstructorArgValue(Arrays.asList(bean));
                if (ctx.getAutowireCapableBeanFactory() instanceof BeanDefinitionRegistry bd) {
                    bd.registerBeanDefinition("MethodToolCallbackProvider-ToolAutoReg" + c.addAndGet(1), newBean.getBeanDefinition());
                }
                break;
            }
        }
        return bean;
    }

    ApplicationContext ctx;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        ctx = applicationContext;
    }

}
